from rest_framework.renderers import JSONRenderer


class CustomRenderer(JSONRenderer):
    def _get_message_and_code(self, data, renderer_context):
        if isinstance(data, dict):
            msg = data.pop('message', '请求成功')
            code = data.pop('code', renderer_context["response"].status_code)
        else:
            msg = '请求成功'
            code = renderer_context["response"].status_code
        return msg, code

    def _get_data(self, data):
        if isinstance(data, list):
            return {'list': data, 'total': len(data)}
        return data

    def render(self, data, accepted_media_type=None, renderer_context=None):
        if renderer_context:
            msg, code = self._get_message_and_code(data, renderer_context)
            data = self._get_data(data)
            ret = {
                'msg': msg,
                'code': code,
                'data': data,
            }
            return super().render(ret, accepted_media_type, renderer_context)
        else:
            return super().render(data, accepted_media_type, renderer_context)